#[longjie.long@aiswu99 tools]$ cat draw_loss.py 
#usage: python /mnt/ssd/caict/tools/draw_loss.py caict-u99-llama-2-7b-DS-FT-10-09-14-11.log --save=caict-u99-llama-2-7b-DS-FT-10-09-14-11-loss.png
import sys
import matplotlib.pyplot as plt
import numpy as np

def extract_loss_values(log_file_path):
    loss_values = []
    with open(log_file_path, 'r') as file:
        for line in file:
            if 'lm loss' in line.lower():
                parts = line.split()
                for i, part in enumerate(parts):
                    if 'lm' in part.lower() and 'loss' in parts[i+1].lower():
                        try:
                            loss_value = float(parts[i+2])
                            loss_values.append(loss_value)
                        except (ValueError, IndexError):
                            continue
    return loss_values

def smooth_loss_curve(loss_values, smoothing_factor=0.9):
    smoothed_values = []
    last_value = loss_values[0]  # Initial value
    for value in loss_values:
        smoothed_value = last_value * smoothing_factor + (1 - smoothing_factor) * value
        smoothed_values.append(smoothed_value)
        last_value = smoothed_value
    return smoothed_values

def plot_loss_curve(loss_values, smooth=False, smoothing_factor=0.9, save_path=None):
    if smooth:
        loss_values = smooth_loss_curve(loss_values, smoothing_factor)
    
    plt.figure(figsize=(10, 6))
    plt.plot(loss_values, label='LM Loss')
    plt.xlabel('Iterations')
    plt.ylabel('LM Loss')
    plt.title('LM Loss Curve')
    plt.legend()
    plt.grid(True)
    
    if save_path:
        plt.savefig(save_path)
        print(f"Loss curve saved as {save_path}")
    else:
        plt.show()

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python plot_loss_curve.py <log_file_path> [--smooth] [--smoothing_factor=<factor>] [--save=<output_image_path>]")
        sys.exit(1)
    
    log_file_path = sys.argv[1]
    smooth = '--smooth' in sys.argv
    smoothing_factor = 0.9
    save_path = None

    # Check for custom smoothing factor and save path
    for arg in sys.argv:
        if arg.startswith('--smoothing_factor='):
            try:
                smoothing_factor = float(arg.split('=')[1])
            except ValueError:
                print("Invalid smoothing factor provided. Using default value of 0.9.")
        elif arg.startswith('--save='):
            save_path = arg.split('=')[1]

    loss_values = extract_loss_values(log_file_path)
    
    if loss_values:
        plot_loss_curve(loss_values, smooth=smooth, smoothing_factor=smoothing_factor, save_path=save_path)
    else:
        print("No loss values found in the log file.")

